# -*- coding: utf-8 -*-
"""
# @FileName:     ddpm.py
# @AuthorName:   Sanqi Lu (Lingwei Dang)
# @Institution:  SCUT, Guangzhou, China
# @EmailAddress: lenvondang@163.com
# @CreateTime:   2024/12/22 12:33
"""
import torch
from tqdm import tqdm

class DDPM_x0():

    # n_steps 就是论文里的 T
    def __init__(self,
                 device,
                 n_steps: int,
                 min_beta: float = 0.0001,
                 max_beta: float = 0.02):
        betas = torch.linspace(min_beta, max_beta, n_steps).to(device)
        alphas = 1 - betas
        alpha_bars = torch.empty_like(alphas)
        product = 1
        for i, alpha in enumerate(alphas):
            product *= alpha
            alpha_bars[i] = product
        self.betas = betas
        self.n_steps = n_steps
        self.alphas = alphas
        self.alpha_bars = alpha_bars

    def sample_forward(self, x, t, eps=None):
        alpha_bar = self.alpha_bars[t].reshape(-1, 1, 1, 1)
        alpha_t_weight = 1/ (1-alpha_bar)
        if eps is None:
            eps = torch.randn_like(x)
        res = eps * torch.sqrt(1 - alpha_bar) + torch.sqrt(alpha_bar) * x
        return res, alpha_t_weight

    def sample_x0(self, img_shape, net, device, simple_var=True):
        x_t = torch.randn(img_shape).to(device)
        net = net.to(device)

        with torch.no_grad():
            for tid in tqdm(range(self.n_steps, 0, -1), total=self.n_steps, desc="ddpm_sample_x0_prediction"):
                # (1) 根据 t, x_t 先用模型预测 x0; 
                t_tensor = torch.tensor([tid] * img_shape[0], dtype=torch.long).to(x_t.device)[..., None]
                pred_x0 = net(x_t, t_tensor)

                # (2) 用公式 q(x_t-1|x0, xt) 计算 x_t-1，主要是要计算出 z_t, mu_t, sigma_t
                # 计算 alpha_t 和 alpha_bar_t
                alpha_t = self.alphas[tid-1]
                alpha_bar_t = self.alpha_bars[tid-1]
                betas_t = self.betas[tid-1]
                
                # 计算均值
                if tid > 1:
                    alpha_bar_t_prev = self.alpha_bars[tid-2]
                    mean = (torch.sqrt(alpha_t) * (1 - alpha_bar_t_prev) / (1 - alpha_bar_t)) * x_t + \
                    (torch.sqrt(alpha_bar_t_prev) * betas_t / (1 - alpha_bar_t)) * pred_x0
                        # 计算方差
                    variance = (1 - alpha_bar_t_prev) / (1 - alpha_bar_t) * betas_t
                        # 采样噪声
                    noise = torch.randn_like(x_t)
                    # 更新 x_t
                    x_t = mean + torch.sqrt(variance) * noise
                else:
                    mean = pred_x0
                    variance = 0
                    # 更新 x_t
                    x_t = mean
        
        return x_t